# Copyright 2022 Dechin Chen
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

hnames = {'ACE': {'CH3': ['H1', 'H2', 'H3']},
          'ALA': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB1', 'HB2', 'HB3']},
          'ARG': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'CD': ['HD2', 'HD3'],
                  'NE': ['HE'], 'NH1': ['HH11', 'HH12'], 'NH2': ['HH21', 'HH22']},
          'ASN': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'ND2': ['HD21', 'HD22']},
          'ASP': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3']},
          'CALA': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB1', 'HB2', 'HB3'], 'C': ['OXT']},
          'CARG': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'CD': ['HD2', 'HD3'],
                   'NE': ['HE'], 'NH1': ['HH11', 'HH12'], 'NH2': ['HH21', 'HH22'], 'C': ['OXT']},
          'CASN': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'ND2': ['HD21', 'HD22'], 'C': ['OXT']},
          'CASP': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'C': ['OXT']},
          'CCYS': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'SG': ['HG'], 'C': ['OXT']},
          'CGLN': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'NE2': ['HE21', 'HE22'],
                   'C': ['OXT']},
          'CGLU': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'C': ['OXT']},
          'CGLY': {'N': ['H'], 'CA': ['HA2', 'HA3'], 'C': ['OXT']},
          'CHID': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'ND1': ['HD1'], 'CE1': ['HE1'], 'CD2': ['HD2'],
                   'C': ['OXT']},
          'CHIS': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CE1': ['HE1'], 'NE2': ['HE2'], 'CD2': ['HD2'],
                   'C': ['OXT']},
          'CILE': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB'], 'CG2': ['HG21', 'HG22', 'HG23'], 'CG1': ['HG12', 'HG13'],
                   'CD1': ['HD11', 'HD12', 'HD13'], 'C': ['OXT']},
          'CLEU': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG'], 'CD1': ['HD11', 'HD12', 'HD13'],
                   'CD2': ['HD21', 'HD22', 'HD23'], 'C': ['OXT']},
          'CLYS': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'CD': ['HD2', 'HD3'],
                   'CE': ['HE2', 'HE3'], 'NZ': ['HZ1', 'HZ2', 'HZ3'], 'C': ['OXT']},
          'CMET': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'CE': ['HE1', 'HE2', 'HE3'],
                   'C': ['OXT']},
          'CPHE': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CD1': ['HD1'], 'CE1': ['HE1'], 'CZ': ['HZ'],
                   'CE2': ['HE2'], 'CD2': ['HD2'], 'C': ['OXT']},
          'CPRO': {'CD': ['HD2', 'HD3'], 'CG': ['HG2', 'HG3'], 'CB': ['HB2', 'HB3'], 'CA': ['HA'], 'C': ['OXT']},
          'CSER': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'OG': ['HG'], 'C': ['OXT']},
          'CTHR': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB'], 'CG2': ['HG21', 'HG22', 'HG23'], 'OG1': ['HG1'],
                   'C': ['OXT']},
          'CTRP': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CD1': ['HD1'], 'NE1': ['HE1'], 'CZ2': ['HZ2'],
                   'CH2': ['HH2'], 'CZ3': ['HZ3'], 'CE3': ['HE3'], 'C': ['OXT']},
          'CTYR': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CD1': ['HD1'], 'CE1': ['HE1'], 'OH': ['HH'],
                   'CE2': ['HE2'], 'CD2': ['HD2'], 'C': ['OXT']},
          'CVAL': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB'], 'CG1': ['HG11', 'HG12', 'HG13'], 'CG2': ['HG21', 'HG22',
                                                                                                    'HG23'],
                   'C': ['OXT']},
          'CYS': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'SG': ['HG']},
          'GLN': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'NE2': ['HE21', 'HE22']},
          'GLU': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3']},
          'GLY': {'N': ['H'], 'CA': ['HA2', 'HA3']},
          'HID': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'ND1': ['HD1'], 'CE1': ['HE1'], 'CD2': ['HD2']},
          'HIS': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CE1': ['HE1'], 'NE2': ['HE2'], 'CD2': ['HD2']},
          'ILE': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB'], 'CG2': ['HG21', 'HG22', 'HG23'], 'CG1': ['HG12', 'HG13'],
                  'CD1': ['HD11', 'HD12', 'HD13']},
          'LEU': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG'], 'CD1': ['HD11', 'HD12', 'HD13'],
                  'CD2': ['HD21', 'HD22', 'HD23']},
          'LYS': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'CD': ['HD2', 'HD3'],
                  'CE': ['HE2', 'HE3'], 'NZ': ['HZ1', 'HZ2', 'HZ3']},
          'MET': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'CE': ['HE1', 'HE2', 'HE3']},
          'NALA': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB1', 'HB2', 'HB3']},
          'NARG': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'CD': ['HD2',
                                                                                                             'HD3'],
                   'NE': ['HE'], 'NH1': ['HH11', 'HH12'], 'NH2': ['HH21', 'HH22']},
          'NASN': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'ND2': ['HD21', 'HD22']},
          'NASP': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3']},
          'NCYS': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'SG': ['HG']},
          'NGLN': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'], 'NE2': ['HE21',
                                                                                                              'HE22']},
          'NGLU': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3']},
          'NGLY': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA2', 'HA3']},
          'NHID': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'ND1': ['HD1'], 'CE1': ['HE1'],
                   'CD2': ['HD2']},
          'NHIS': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CE1': ['HE1'], 'NE2': ['HE2'],
                   'CD2': ['HD2']},
          'NILE': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB'], 'CG2': ['HG21', 'HG22', 'HG23'],
                   'CG1': ['HG12', 'HG13'], 'CD1': ['HD11', 'HD12', 'HD13']},
          'NLEU': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG'],
                   'CD1': ['HD11', 'HD12', 'HD13'], 'CD2': ['HD21', 'HD22', 'HD23']},
          'NLYS': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'],
                   'CD': ['HD2', 'HD3'], 'CE': ['HE2', 'HE3'], 'NZ': ['HZ1', 'HZ2', 'HZ3']},
          'NME': {'N': ['H'], 'CH3': ['HH31', 'HH32', 'HH33']},
          'NMET': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CG': ['HG2', 'HG3'],
                   'CE': ['HE1', 'HE2', 'HE3']},
          'NPHE': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CD1': ['HD1'], 'CE1': ['HE1'],
                   'CZ': ['HZ'], 'CE2': ['HE2'], 'CD2': ['HD2']},
          'NPRO': {'N': ['H2', 'H3'], 'CD': ['HD2', 'HD3'], 'CG': ['HG2', 'HG3'], 'CB': ['HB2', 'HB3'], 'CA': ['HA']},
          'NSER': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'OG': ['HG']},
          'NTHR': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB'], 'CG2': ['HG21', 'HG22', 'HG23'],
                   'OG1': ['HG1']},
          'NTRP': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CD1': ['HD1'], 'NE1': ['HE1'],
                   'CZ2': ['HZ2'], 'CH2': ['HH2'], 'CZ3': ['HZ3'], 'CE3': ['HE3']},
          'NTYR': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CD1': ['HD1'], 'CE1': ['HE1'],
                   'OH': ['HH'], 'CE2': ['HE2'], 'CD2': ['HD2']},
          'NVAL': {'N': ['H1', 'H2', 'H3'], 'CA': ['HA'], 'CB': ['HB'], 'CG1': ['HG11', 'HG12', 'HG13'],
                   'CG2': ['HG21', 'HG22', 'HG23']},
          'PHE': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CD1': ['HD1'], 'CE1': ['HE1'], 'CZ': ['HZ'],
                  'CE2': ['HE2'], 'CD2': ['HD2']},
          'PRO': {'CD': ['HD2', 'HD3'], 'CG': ['HG2', 'HG3'], 'CB': ['HB2', 'HB3'], 'CA': ['HA']},
          'SER': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'OG': ['HG']},
          'THR': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB'], 'CG2': ['HG21', 'HG22', 'HG23'], 'OG1': ['HG1']},
          'TRP': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CD1': ['HD1'], 'NE1': ['HE1'], 'CZ2': ['HZ2'],
                  'CH2': ['HH2'], 'CZ3': ['HZ3'], 'CE3': ['HE3']},
          'TYR': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB2', 'HB3'], 'CD1': ['HD1'], 'CE1': ['HE1'], 'OH': ['HH'],
                  'CE2': ['HE2'], 'CD2': ['HD2']},
          'VAL': {'N': ['H'], 'CA': ['HA'], 'CB': ['HB'], 'CG1': ['HG11', 'HG12', 'HG13'], 'CG2': ['HG21',
                                                                                                   'HG22', 'HG23']},
          'WAT': {'O': ['H1', 'H2']},
          'HOH': {'O': ['H1', 'H2']}
          }

hbond_type = {
    'ACE': {
        'CH3': np.array(['ch3', 'C', 'O'])
    },
    'ALA': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'CB']),
        'CB': np.array(['ch3', 'CA', 'C'])
    },
    'ARG': {
        'N': np.array(['dihedral', 'CA', 'CB']),
        'CA': np.array(['cc3', 'CB', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'CD': np.array(['c2h2', 'CG', 'NE']),
        'NE': np.array(['c6', 'CD', 'CZ']),
        'NH1': np.array([['dihedral', 'CZ', 'NH2'],
                         ['dihedral', 'CZ', 'NE']]),
        'NH2': np.array([['dihedral', 'CZ', 'NH1'],
                         ['dihedral', 'CZ', 'NE']])
    },
    'ASN': {
        'ND2': np.array(['c2h4', 'CG', 'OD1']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'CB'])
    },
    'ASP': {
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'CB'])
    },
    'CALA': {
        'N': np.array(['dihedral', 'CA', 'CB']),
        'CA': np.array(['cc3', 'CB', 'C']),
        'CB': np.array(['ch3', 'CA', 'C']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CARG': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'CD': np.array(['c2h2', 'CG', 'NE']),
        'NE': np.array(['dihedral', 'CZ', 'NH1']),
        'NH1': np.array([['dihedral', 'CZ', 'NH2'],
                         ['dihedral', 'CZ', 'NE']]),
        'NH2': np.array([['dihedral', 'CZ', 'NH1'],
                         ['dihedral', 'CZ', 'NE']]),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CASN': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'ND2': np.array(['c2h4', 'CG', 'OD1']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CASP': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CCYS': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'SG']),
        'SG': np.array(['dihedral', 'CB', 'CA']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CGLN': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'NE2': np.array(['c2h4', 'CD', 'OE1']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CGLU': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CGLY': {
        'CA': np.array(['c2h2', 'N', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CHID': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD2': np.array(['c6', 'CG', 'NE2']),
        'ND1': np.array(['c6', 'CG', 'CE1']),
        'CE1': np.array(['c6', 'ND1', 'NE2']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CHIS': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD2': np.array(['c6', 'CG', 'NE2']),
        'NE2': np.array(['c6', 'CD2', 'CE1']),
        'CE1': np.array(['c6', 'ND1', 'NE2']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CILE': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['cc3', 'CG1', 'CA']),
        'CG2': np.array(['ch3', 'CB', 'CA']),
        'CG1': np.array(['c2h2', 'CB', 'CD1']),
        'CD1': np.array(['ch3', 'CG1', 'CB']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CLEU': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['cc3', 'CD2', 'CB']),
        'CD2': np.array(['ch3', 'CG', 'CD1']),
        'CD1': np.array(['ch3', 'CG', 'CB']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CLYS': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'CD': np.array(['c2h2', 'CG', 'CE']),
        'CE': np.array(['c2h2', 'CD', 'NZ']),
        'NZ': np.array(['ch3', 'CE', 'CD']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CMET': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'SD']),
        'CE': np.array(['ch3', 'SD', 'CG']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CPHE': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD1': np.array(['c6', 'CG', 'CE1']),
        'CD2': np.array(['c6', 'CG', 'CE2']),
        'CE1': np.array(['c6', 'CD1', 'CZ']),
        'CE2': np.array(['c6', 'CD2', 'CZ']),
        'CZ': np.array(['c6', 'CE2', 'CE1']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CPRO': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'CD': np.array(['c2h2', 'CG', 'N']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CSER': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'OG']),
        'OG': np.array(['dihedral', 'CB', 'CA']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CTHR': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['cc3', 'CG2', 'OG1']),
        'OG1': np.array(['dihedral', 'CB', 'CA']),
        'CG2': np.array(['ch3', 'CB', 'CA']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CTRP': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD1': np.array(['c6', 'CG', 'NE1']),
        'NE1': np.array(['c6', 'CD1', 'CE2']),
        'CZ2': np.array(['c6', 'CE2', 'CH2']),
        'CH2': np.array(['c6', 'CZ2', 'CZ3']),
        'CZ3': np.array(['c6', 'CH2', 'CE3']),
        'CE3': np.array(['c6', 'CD2', 'CZ3']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CTYR': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD1': np.array(['c6', 'CG', 'CE1']),
        'CE1': np.array(['c6', 'CD1', 'CZ']),
        'CD2': np.array(['c6', 'CG', 'CE2']),
        'CE2': np.array(['c6', 'CD2', 'CZ']),
        'OH': np.array(['dihedral', 'CZ', 'CE2']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CVAL': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['cc3', 'CG2', 'CA']),
        'CG1': np.array(['ch3', 'CB', 'CA']),
        'CG2': np.array(['ch3', 'CB', 'CA']),
        'C': np.array(['dihedral', 'CA', 'N'])
    },
    'CYS': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'SG']),
        'SG': np.array(['dihedral', 'CB', 'CA'])
    },
    'GLN': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'NE2': np.array(['c2h4', 'CD', 'OE1'])
    },
    'GLU': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD'])
    },
    'GLY': {
        'CA': np.array(['c2h2', 'N', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
    },
    'HID': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD2': np.array(['c6', 'CG', 'NE2']),
        'ND1': np.array(['c6', 'CG', 'CE1']),
        'CE1': np.array(['c6', 'ND1', 'NE2'])
    },
    'HIS': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD2': np.array(['c6', 'CG', 'NE2']),
        'NE2': np.array(['c6', 'CD2', 'CE1']),
        'CE1': np.array(['c6', 'ND1', 'NE2'])
    },
    'ILE': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['cc3', 'CG1', 'CA']),
        'CG2': np.array(['ch3', 'CB', 'CA']),
        'CG1': np.array(['c2h2', 'CB', 'CD1']),
        'CD1': np.array(['ch3', 'CG1', 'CB'])
    },
    'LEU': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['cc3', 'CD2', 'CB']),
        'CD2': np.array(['ch3', 'CG', 'CD1']),
        'CD1': np.array(['ch3', 'CG', 'CB'])
    },
    'LYS': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'CD': np.array(['c2h2', 'CG', 'CE']),
        'CE': np.array(['c2h2', 'CD', 'NZ']),
        'NZ': np.array(['ch3', 'CE', 'CD'])
    },
    'MET': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'SD']),
        'CE': np.array(['ch3', 'SD', 'CG'])
    },
    'NALA': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'CB']),
        'CB': np.array(['ch3', 'CA', 'C'])
    },
    'NARG': {
        'N': np.array(['ch3', 'CA', 'C']),
        'CA': np.array(['cc3', 'CB', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'CD': np.array(['c2h2', 'CG', 'NE']),
        'NE': np.array(['dihedral', 'CZ', 'NH1']),
        'NH1': np.array([['dihedral', 'CZ', 'NH2'],
                         ['dihedral', 'CZ', 'NE']]),
        'NH2': np.array([['dihedral', 'CZ', 'NH1'],
                         ['dihedral', 'CZ', 'NE']])
    },
    'NASN': {
        'ND2': np.array(['c2h4', 'CG', 'OD1']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C'])
    },
    'NASP': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG'])
    },
    'NCYS': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'SG']),
        'SG': np.array(['dihedral', 'CB', 'CA'])
    },
    'NGLN': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'NE2': np.array(['c2h4', 'CD', 'OE1'])
    },
    'NGLU': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD'])
    },
    'NGLY': {
        'CA': np.array(['c2h2', 'N', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
    },
    'NHID': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD2': np.array(['c6', 'CG', 'NE2']),
        'ND1': np.array(['c6', 'CG', 'CE1']),
        'CE1': np.array(['c6', 'ND1', 'NE2'])
    },
    'NHIS': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD2': np.array(['c6', 'CG', 'NE2']),
        'NE2': np.array(['c6', 'CD2', 'CE1']),
        'CE1': np.array(['c6', 'ND1', 'NE2'])
    },
    'NILE': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['cc3', 'CG1', 'CA']),
        'CG2': np.array(['ch3', 'CB', 'CA']),
        'CG1': np.array(['c2h2', 'CB', 'CD1']),
        'CD1': np.array(['ch3', 'CG1', 'CB'])
    },
    'NLEU': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['cc3', 'CD2', 'CB']),
        'CD2': np.array(['ch3', 'CG', 'CD1']),
        'CD1': np.array(['ch3', 'CG', 'CB'])
    },
    'NLYS': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'CD': np.array(['c2h2', 'CG', 'CE']),
        'CE': np.array(['c2h2', 'CD', 'NZ']),
        'NZ': np.array(['ch3', 'CE', 'CD'])
    },
    'NMET': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'SD']),
        'CE': np.array(['ch3', 'SD', 'CG'])
    },
    'NPHE': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD1': np.array(['c6', 'CG', 'CE1']),
        'CD2': np.array(['c6', 'CG', 'CE2']),
        'CE1': np.array(['c6', 'CD1', 'CZ']),
        'CE2': np.array(['c6', 'CD2', 'CZ']),
        'CZ': np.array(['c6', 'CE2', 'CE1']),
    },
    'NPRO': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'CD': np.array(['c2h2', 'CG', 'N']),
        'N': np.array(['c2h2', 'CA', 'CD'])
    },
    'NSER': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'OG']),
        'OG': np.array(['dihedral', 'CB', 'CA'])
    },
    'NTHR': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['cc3', 'CG2', 'OG1']),
        'OG1': np.array(['dihedral', 'CB', 'CA']),
        'CG2': np.array(['ch3', 'CB', 'CA'])
    },
    'NTRP': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD1': np.array(['c6', 'CG', 'NE1']),
        'NE1': np.array(['c6', 'CD1', 'CE2']),
        'CZ2': np.array(['c6', 'CE2', 'CH2']),
        'CH2': np.array(['c6', 'CZ2', 'CZ3']),
        'CZ3': np.array(['c6', 'CH2', 'CE3']),
        'CE3': np.array(['c6', 'CD2', 'CZ3'])
    },
    'NTYR': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD1': np.array(['c6', 'CG', 'CE1']),
        'CE1': np.array(['c6', 'CD1', 'CZ']),
        'CD2': np.array(['c6', 'CG', 'CE2']),
        'CE2': np.array(['c6', 'CD2', 'CZ']),
        'OH': np.array(['dihedral', 'CZ', 'CE2'])
    },
    'NVAL': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['ch3', 'CA', 'C']),
        'CB': np.array(['cc3', 'CG2', 'CA']),
        'CG1': np.array(['ch3', 'CB', 'CA']),
        'CG2': np.array(['ch3', 'CB', 'CA'])
    },
    'PHE': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD1': np.array(['c6', 'CG', 'CE1']),
        'CD2': np.array(['c6', 'CG', 'CE2']),
        'CE1': np.array(['c6', 'CD1', 'CZ']),
        'CE2': np.array(['c6', 'CD2', 'CZ']),
        'CZ': np.array(['c6', 'CE2', 'CE1']),
    },
    'PRO': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CG': np.array(['c2h2', 'CB', 'CD']),
        'CD': np.array(['c2h2', 'CG', 'N']),
    },
    'SER': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'OG']),
        'OG': np.array(['dihedral', 'CB', 'CA'])
    },
    'THR': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['cc3', 'CG2', 'OG1']),
        'OG1': np.array(['dihedral', 'CB', 'CA']),
        'CG2': np.array(['ch3', 'CB', 'CA'])
    },
    'TRP': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD1': np.array(['c6', 'CG', 'NE1']),
        'NE1': np.array(['c6', 'CD1', 'CE2']),
        'CZ2': np.array(['c6', 'CE2', 'CH2']),
        'CH2': np.array(['c6', 'CZ2', 'CZ3']),
        'CZ3': np.array(['c6', 'CH2', 'CE3']),
        'CE3': np.array(['c6', 'CD2', 'CZ3'])
    },
    'TYR': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['c2h2', 'CA', 'CG']),
        'CD1': np.array(['c6', 'CG', 'CE1']),
        'CE1': np.array(['c6', 'CD1', 'CZ']),
        'CD2': np.array(['c6', 'CG', 'CE2']),
        'CE2': np.array(['c6', 'CD2', 'CZ']),
        'OH': np.array(['dihedral', 'CZ', 'CE2'])
    },
    'VAL': {
        'CA': np.array(['cc3', 'CB', 'C']),
        'N': np.array(['dihedral', 'CA', 'C']),
        'CB': np.array(['cc3', 'CG2', 'CA']),
        'CG1': np.array(['ch3', 'CB', 'CA']),
        'CG2': np.array(['ch3', 'CB', 'CA'])
    },
    'NME': {
        'N': np.array(['c6', 'CH3', 'C']),
        'CH3': np.array(['ch3', 'N', 'C'])
    },
    'WAT': {
        'O': np.array(['wat', 'O', 'O'])
    },
    'HOH': {
        'O': np.array(['wat', 'O', 'O'])
    }
}

restypes = [
    'A', 'R', 'N', 'D', 'C',
    'Q', 'E', 'G', 'H', 'I',
    'L', 'K', 'M', 'F', 'P',
    'S', 'T', 'W', 'Y', 'V'
]

resdict = {'ALA': 0, 'ARG': 1, 'ASN': 2, 'ASP': 3, 'CYS': 4,
           'GLN': 5, 'GLU': 6, 'GLY': 7, 'HIS': 8, 'ILE': 9, 'HID': 8,
           'LEU': 10, 'LYS': 11, 'MET': 12, 'PHE': 13, 'PRO': 14,
           'SER': 15, 'THR': 16, 'TRP': 17, 'TYR': 18, 'VAL': 19,
           'CALA': 0, 'CARG': 1, 'CASN': 2, 'CASP': 3, 'CCYS': 4,
           'CGLN': 5, 'CGLU': 6, 'CGLY': 7, 'CHIS': 8, 'CILE': 9, 'CHID': 8,
           'CLEU': 10, 'CLYS': 11, 'CMET': 12, 'CPHE': 13, 'CPRO': 14,
           'CSER': 15, 'CTHR': 16, 'CTRP': 17, 'CTYR': 18, 'CVAL': 19,
           'NALA': 0, 'NARG': 1, 'NASN': 2, 'NASP': 3, 'NCYS': 4,
           'NGLN': 5, 'NGLU': 6, 'NGLY': 7, 'NHIS': 8, 'NILE': 9, 'NHID': 8,
           'NLEU': 10, 'NLYS': 11, 'NMET': 12, 'NPHE': 13, 'NPRO': 14,
           'NSER': 15, 'NTHR': 16, 'NTRP': 17, 'NTYR': 18, 'NVAL': 19,
           'CHIE': 8, 'HIE': 8, 'NHIE': 8, 'ACE': 21, 'NME': 22, 'WAT': 23, 'HOH': 23
           }

atom_types = [
    'N', 'CA', 'C', 'CB', 'O', 'CG', 'CG1', 'CG2', 'OG', 'OG1', 'SG', 'CD',
    'CD1', 'CD2', 'ND1', 'ND2', 'OD1', 'OD2', 'SD', 'CE', 'CE1', 'CE2', 'CE3',
    'NE', 'NE1', 'NE2', 'OE1', 'OE2', 'CH2', 'NH1', 'NH2', 'OH', 'CZ', 'CZ2',
    'CZ3', 'NZ', 'OXT', 'CH3'
]

restype_name_to_atom14_names = {
    'ALA': ['N', 'CA', 'C', 'O', 'CB', '', '', '', '', '', '', '', '', ''],
    'ARG': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', 'NE', 'CZ', 'NH1', 'NH2', '', '', ''],
    'ASN': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'OD1', 'ND2', '', '', '', '', '', ''],
    'ASP': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'OD1', 'OD2', '', '', '', '', '', ''],
    'CYS': ['N', 'CA', 'C', 'O', 'CB', 'SG', '', '', '', '', '', '', '', ''],
    'GLN': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', 'OE1', 'NE2', '', '', '', '', ''],
    'GLU': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', 'OE1', 'OE2', '', '', '', '', ''],
    'GLY': ['N', 'CA', 'C', 'O', '', '', '', '', '', '', '', '', '', ''],
    'HIS': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'ND1', 'CD2', 'CE1', 'NE2', '', '', '', ''],
    'HIE': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'ND1', 'CD2', 'CE1', 'NE2', '', '', '', ''],
    'HID': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'ND1', 'CD2', 'CE1', 'NE2', '', '', '', ''],
    'ILE': ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', 'CD1', '', '', '', '', '', ''],
    'LEU': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD1', 'CD2', '', '', '', '', '', ''],
    'LYS': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', 'CE', 'NZ', '', '', '', '', ''],
    'MET': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'SD', 'CE', '', '', '', '', '', ''],
    'PHE': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', '', '', ''],
    'PRO': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', '', '', '', '', '', '', ''],
    'SER': ['N', 'CA', 'C', 'O', 'CB', 'OG', '', '', '', '', '', '', '', ''],
    'THR': ['N', 'CA', 'C', 'O', 'CB', 'OG1', 'CG2', '', '', '', '', '', '', ''],
    'TRP': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD1', 'CD2', 'NE1', 'CE2', 'CE3', 'CZ2', 'CZ3', 'CH2'],
    'TYR': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', 'OH', '', ''],
    'VAL': ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', '', '', '', '', '', '', ''],
    'UNK': ['', '', '', '', '', '', '', '', '', '', '', '', '', ''],
}

restype_name_to_atom14_masks = {
    'ALA': [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    'ARG': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
    'ASN': [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
    'ASP': [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
    'CYS': [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
    'GLN': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
    'GLU': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
    'GLY': [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    'HIS': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
    'HIE': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
    'HID': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
    'ILE': [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
    'LEU': [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
    'LYS': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
    'MET': [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
    'PHE': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
    'PRO': [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
    'SER': [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
    'THR': [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
    'TRP': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
    'TYR': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
    'VAL': [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
    'UNK': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]}

atom14_order_dict = {'ACE': {'CH3': 0, 'C': 1, 'O': 2},
                     'ALA': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4},
                     'ARG': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'CD': 6,
                             'NE': 7,
                             'CZ': 8,
                             'NH1': 9,
                             'NH2': 10},
                     'ASN': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'OD1': 6,
                             'ND2': 7},
                     'ASP': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'OD1': 6,
                             'OD2': 7},
                     'CYS': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4, 'SG': 5},
                     'GLN': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'CD': 6,
                             'OE1': 7,
                             'NE2': 8},
                     'GLU': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'CD': 6,
                             'OE1': 7,
                             'OE2': 8},
                     'GLY': {'N': 0, 'CA': 1, 'C': 2, 'O': 3},
                     'HIS': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'ND1': 6,
                             'CD2': 7,
                             'CE1': 8,
                             'NE2': 9},
                     'HIE': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'ND1': 6,
                             'CD2': 7,
                             'CE1': 8,
                             'NE2': 9},
                     'HID': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'ND1': 6,
                             'CD2': 7,
                             'CE1': 8,
                             'NE2': 9},
                     'ILE': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG1': 5,
                             'CG2': 6,
                             'CD1': 7},
                     'LEU': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'CD1': 6,
                             'CD2': 7},
                     'LYS': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'CD': 6,
                             'CE': 7,
                             'NZ': 8},
                     'MET': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4, 'CG': 5, 'SD': 6, 'CE': 7},
                     'PHE': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'CD1': 6,
                             'CD2': 7,
                             'CE1': 8,
                             'CE2': 9,
                             'CZ': 10},
                     'PRO': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4, 'CG': 5, 'CD': 6},
                     'SER': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4, 'OG': 5},
                     'THR': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4, 'OG1': 5, 'CG2': 6},
                     'TRP': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'CD1': 6,
                             'CD2': 7,
                             'NE1': 8,
                             'CE2': 9,
                             'CE3': 10,
                             'CZ2': 11,
                             'CZ3': 12,
                             'CH2': 13},
                     'TYR': {'N': 0,
                             'CA': 1,
                             'C': 2,
                             'O': 3,
                             'CB': 4,
                             'CG': 5,
                             'CD1': 6,
                             'CD2': 7,
                             'CE1': 8,
                             'CE2': 9,
                             'CZ': 10,
                             'OH': 11},
                     'VAL': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4, 'CG1': 5, 'CG2': 6},
                     'NME': {'N': 0, 'CH3': 1},
                     'UNK': {}}

atom14_to_atom37_dict = {'ALA': [0, 1, 2, 4, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                         'ARG': [0, 1, 2, 4, 3, 5, 11, 23, 32, 29, 30, 0, 0, 0],
                         'ASN': [0, 1, 2, 4, 3, 5, 16, 15, 0, 0, 0, 0, 0, 0],
                         'ASP': [0, 1, 2, 4, 3, 5, 16, 17, 0, 0, 0, 0, 0, 0],
                         'CYS': [0, 1, 2, 4, 3, 10, 0, 0, 0, 0, 0, 0, 0, 0],
                         'GLN': [0, 1, 2, 4, 3, 5, 11, 26, 25, 0, 0, 0, 0, 0],
                         'GLU': [0, 1, 2, 4, 3, 5, 11, 26, 27, 0, 0, 0, 0, 0],
                         'GLY': [0, 1, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                         'HIS': [0, 1, 2, 4, 3, 5, 14, 13, 20, 25, 0, 0, 0, 0],
                         'HIE': [0, 1, 2, 4, 3, 5, 14, 13, 20, 25, 0, 0, 0, 0],
                         'HID': [0, 1, 2, 4, 3, 5, 14, 13, 20, 25, 0, 0, 0, 0],
                         'ILE': [0, 1, 2, 4, 3, 6, 7, 12, 0, 0, 0, 0, 0, 0],
                         'LEU': [0, 1, 2, 4, 3, 5, 12, 13, 0, 0, 0, 0, 0, 0],
                         'LYS': [0, 1, 2, 4, 3, 5, 11, 19, 35, 0, 0, 0, 0, 0],
                         'MET': [0, 1, 2, 4, 3, 5, 18, 19, 0, 0, 0, 0, 0, 0],
                         'PHE': [0, 1, 2, 4, 3, 5, 12, 13, 20, 21, 32, 0, 0, 0],
                         'PRO': [0, 1, 2, 4, 3, 5, 11, 0, 0, 0, 0, 0, 0, 0],
                         'SER': [0, 1, 2, 4, 3, 8, 0, 0, 0, 0, 0, 0, 0, 0],
                         'THR': [0, 1, 2, 4, 3, 9, 7, 0, 0, 0, 0, 0, 0, 0],
                         'TRP': [0, 1, 2, 4, 3, 5, 12, 13, 24, 21, 22, 33, 34, 28],
                         'TYR': [0, 1, 2, 4, 3, 5, 12, 13, 20, 21, 32, 31, 0, 0],
                         'VAL': [0, 1, 2, 4, 3, 6, 7, 0, 0, 0, 0, 0, 0, 0],
                         'UNK': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]}

RESIDUE_NAMES = np.array(['ALA', 'ARG', 'ASN', 'ASP', 'CYS', 'GLN', 'GLU', 'GLY', 'HID', 'HIS',
                          'ILE', 'LEU', 'LYS', 'MET', 'PHE', 'PRO', 'SER', 'THR', 'TRP', 'TYR', 'VAL'], np.str_)
